{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  WumpusExpr.KNormal
-- Copyright   :  (c) Stephen Tetley 2012
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- KNormal form: @ SyntaxExp -> KNormal @
--
--------------------------------------------------------------------------------

module WumpusExpr.KNormal
  (

    Program(..)
  , Fundef(..)
  , Expr(..)
  , Command(..)
  , Value(..)
  , typeOf
  , typeOfV
  
  , ppProgram
 
  , FreeVars
  , freevars
  , knormal

  ) where

import WumpusExpr.FreshMon
import WumpusExpr.Operators
import WumpusExpr.Radian
import qualified WumpusExpr.Syntax as S
import WumpusExpr.Types
import WumpusExpr.Utils
import WumpusExpr.VarId

import Text.PrettyExprHPJ                   -- package: pretty-expr-hpj
import Text.PrettyExprHPJ.Haskell98

import Text.PrettyPrint.HughesPJ            -- package: pretty

import Control.Applicative
import qualified Data.Set as S


data Program = Program 
      { proc_main :: Fundef 
      , proc_defs :: [Fundef]
      }
  deriving (Eq,Show)


data Fundef = Fundef 
    { fun_name       :: VarId
    , fun_args       :: [VarId]
    , fun_body       :: Expr
    } 
  deriving (Eq, Show)


-- | Prim is prim-application.
--
data Expr = Val Value
          | FunCall VarId [Value]
          | Let VarId Expr Expr
          | Ifte Value Expr Expr
          | Cmd Command
          | Expr :>> Expr
          | Return Expr
  deriving (Eq,Show)


-- | Separation of Values (and simple expressions) and 
-- complex Expressions.
--
-- Including unary, binary and relation expressions as simple
-- expressions makes constant folding and useless conditional
-- elimination easier.
-- 
data Value = Unit
           | Bool Bool
           | Int Int
           | Float Decimal
           | Angle Radian
           | String String
           | PsSymbol String
           | Var VarId
           | UnaryV UnaryOp Value
           | BinV BinOp Value Value
           | RelV RelOp Value Value
  deriving (Eq,Show)


data Command = NewPath
             | ClosePath
             | Fill
             | Stroke
             | MoveTo Value Value
             | LineTo Value Value
             | CurveTo Value Value Value Value Value Value
             | GSave
             | GRestore
             | Scale Value Value
             | Rotate Value
             | Translate Value Value
             | Transform Value Value Value Value Value Value
             | ShowText Value
  deriving (Eq,Show)



-- | Note the type of a FunCall is its answer type
-- (i.e. the application of Prim to its arguments) not the type 
-- of the opcode function.
--
-- C.f. Haskell
--
-- > :t ( 1 + 2)  == Int  (Prim App)
-- 
-- > :t (+)       == Int -> Int -> Int (Prim Def)
--
--
typeOf :: Expr -> Type
typeOf (Val v)                  = typeOfV v

typeOf (FunCall v _)            = 
    case varType v of { TyFun _ a -> a; t -> t }

typeOf (Let _ _ e)              = typeOf e
typeOf (Ifte _  _ e2)           = typeOf e2
typeOf (Cmd {})                 = TyUnit
typeOf (_ :>> e2)               = TyMon $ typeOf e2
typeOf (Return e)               = TyMon $ typeOf e


typeOfV :: Value -> Type
typeOfV (Unit)                   = TyUnit
typeOfV (Bool {})                = TyBool
typeOfV (Int {})                 = TyInt
typeOfV (Float {})               = TyFloat
typeOfV (Angle {})               = TyRadian
typeOfV (String {})              = TyString
typeOfV (PsSymbol {})            = TyPsSymbol
typeOfV (Var v)                  = varType v
typeOfV (UnaryV _ v)             = typeOfV v
typeOfV (BinV _ v1 _)            = typeOfV v1
typeOfV (RelV {})                = TyBool



--------------------------------------------------------------------------------
-- pretty print

ppProgram :: Program -> Doc
ppProgram (Program main_proc defs) = 
    vcatBlank $ map ppFundef (main_proc : defs)


ppFundef :: Fundef -> Doc
ppFundef (Fundef hd args e) = 
    text "def" <+> ppVarId hd <+> hsep (map ppVarId args) <+> char '=' 
      $+$ nest 2 (ppExpr e)

ppExpr :: Expr -> Doc
ppExpr = unparse . exprDoc



exprDoc :: Expr -> DocE
exprDoc (Val v)                     = exprValue v

exprDoc (FunCall v es)              = 
    funAppB (literal $ varName v) (map exprValue es)

exprDoc (Let vid e1 e2)             = 
    let pre = text "let" <+> text (varName vid) 
    in Atom $ letBody pre e1 e2

exprDoc (Ifte v1 e1 e2)    = 
    Atom $ text "if" <+> unparse (exprValue v1)
                     <+> nest 2 (    text "then" <+> ppExpr e1
                                 $+$ text "else" <+> ppExpr e2)

exprDoc (Cmd cmd)                   = ppCommandE cmd

exprDoc (e1 :>> e2)                 = 
    Atom $ ppExpr e1 <> char ';' $+$ ppExpr e2

exprDoc (Return e)                  = funAppB1 (literal "return") (exprDoc e)

ppCommandE :: Command -> DocE
ppCommandE (NewPath)                    = funAppB (literal "newpath") []
ppCommandE (ClosePath)                  = funAppB (literal "closepath") []
ppCommandE (Fill)                       = funAppB (literal "fill") []
ppCommandE (Stroke)                     = funAppB (literal "stroke") []
ppCommandE (MoveTo x y)                 = 
    funAppB (literal "moveto") (map exprValue [x,y])

ppCommandE (LineTo x y)                 = 
    funAppB (literal "lineto") (map exprValue [x,y])

ppCommandE (CurveTo x1 y1 x2 y2 x3 y3)  =
    funAppB (literal "curveto") (map exprValue [x1,y1,x2,y2,x3,y3])

ppCommandE (GSave)                      = funAppB (literal "gsave") []
ppCommandE (GRestore)                   = funAppB (literal "grestore") []
ppCommandE (Scale x y)                  = 
    funAppB (literal "scale") (map exprValue [x,y])

ppCommandE (Rotate ang)                 = 
    funAppB (literal "rotate") [exprValue ang]

ppCommandE (Translate x y)              = 
    funAppB (literal "translate") (map exprValue [x,y])

ppCommandE (Transform a b c d e f)      = 
    funAppB (literal "transform") (map exprValue [a,b,c,d,e,f])

ppCommandE (ShowText v)                 = 
    funAppB (literal "scale") [exprValue v]





exprValue :: Value -> DocE
exprValue (Unit)                      = literal "()"
exprValue (Bool b)                    = 
    Atom $ text $ if b then "True" else "False"

exprValue (Int i)                     = Atom $ int i
exprValue (Float d)                   = Atom $ truncatedDouble $ realToFrac d
exprValue (Angle d)                   = Atom $ truncatedDouble $ realToFrac d
exprValue (String ss)                 = Atom $ doubleQuotes $ text ss
exprValue (PsSymbol ss)               = Atom $ text ('/':ss)
exprValue (Var v)                     = literal $ varName v 
exprValue (UnaryV op v)               = unaryDocE op (exprValue v)
exprValue (BinV op v1 v2)             = binDocE op (exprValue v1) (exprValue v2)
exprValue (RelV op v1 v2)             = relDocE op (exprValue v1) (exprValue v2)

letBody :: Doc -> Expr -> Expr -> Doc
letBody d e b = 
    d <+> char '=' <+> ppExpr e <+> text "in" $+$ nest 2 (ppExpr b)



--------------------------------------------------------------------------------
-- Free vars

type FreeVars = S.Set VarId

freevars :: Expr -> FreeVars
freevars (Val v)                = freevarsV v
freevars (Let v e1 e2)          =
    S.union (freevars e1) (S.delete v $ freevars e2)

freevars (FunCall _ xs)         = S.unions $ map freevarsV xs


freevars (Ifte v1 e1 e2)   = 
    let s1 = freevarsV v1 in S.unions [s1, freevars e1, freevars e2]

freevars (Cmd cmd)              = freevarsC cmd

freevars (e1 :>> e2)            = freevars e1 `S.union` freevars e2
freevars (Return e)             = freevars e

freevarsV :: Value -> FreeVars
freevarsV (Var v)                = S.singleton v
freevarsV (UnaryV _ v)           = freevarsV v
freevarsV (BinV _ v1 v2)         = freevarsV v1 `S.union` freevarsV v2
freevarsV (RelV _ v1 v2)         = freevarsV v1 `S.union` freevarsV v2
freevarsV _                      = S.empty


freevarsC :: Command -> FreeVars
freevarsC (MoveTo x y)                  = freevarsVs [x,y]
freevarsC (LineTo x y)                  = freevarsVs [x,y]
freevarsC (CurveTo x1 y1 x2 y2 x3 y3)   = freevarsVs [x1,y1,x2,y2,x3,y3]
freevarsC (Scale x y)                   = freevarsVs [x,y]
freevarsC (Rotate ang)                  = freevarsVs [ang]
freevarsC (Translate x y)               = freevarsVs [x,y]
freevarsC (Transform a b c d e f)       = freevarsVs [a,b,c,d,e,f]
freevarsC (ShowText v)                  = freevarsVs [v]
freevarsC _                             = S.empty

freevarsVs :: [Value] -> FreeVars
freevarsVs = S.unions . map freevarsV


--------------------------------------------------------------------------------
-- K Normalization


infixl 1 `normalizeToLet`

normalizeToLet :: S.Expr
               -> (Value -> Mon Expr)
               -> Mon Expr
normalizeToLet e0 k = knormalE e0 >>= \a -> case a of
    (Val v)       -> k v
    e             -> do { vnew  <- fresh (typeOf e)
                        ; e'    <- k (Var vnew)
                        ; return $ Let vnew e e'
                        }


-- does knormalization need bindings if VarIds carry their type?


infixl 1 `normalizeToLets`

normalizeToLets :: [S.Expr]
                -> ([Value] -> Mon Expr)
                -> Mon Expr
normalizeToLets es0 k = go emptyH es0
  where
    go h []     = k (fromH h)
    go h (e:es) = normalizeToLet e $ \v -> go (snocH h v) es


knormal :: S.Program -> Mon Program
knormal (S.Program main_proc defs) =  
    Program <$> knormalF main_proc <*> mapM knormalF defs


knormalF :: S.Fundef -> Mon Fundef
knormalF (S.Fundef hd args body)  = Fundef hd args <$> knormalE body

--
-- Note - for fun application we need a bit more care 
-- vis-a-vis names and alpha renaming...
--

knormalE :: S.Expr -> Mon Expr
knormalE (S.Unit)               = return $ Val Unit
knormalE (S.Bool b)             = return $ Val $ Int $ if b then 1 else 0
knormalE (S.Int i)              = return $ Val $ Int i
knormalE (S.Float d)            = return $ Val $ Float d
knormalE (S.Angle d)            = return $ Val $ Angle d
knormalE (S.String ss)          = return $ Val $ String ss
knormalE (S.PsSymbol ss)        = return $ Val $ PsSymbol ss
knormalE (S.Var v)              = return $ Val $ Var v
knormalE (S.App e es)           = 
    e  `normalizeToLet`   \v1   -> 
    es `normalizeToLets`  \vals -> 
    case v1 of 
      (Var v) -> return $ FunCall v vals
      _       -> throwError $ "Bad Bpp..."


knormalE (S.Let vid e1 e2)      = do 
    ea <- knormalE e1
    eb <- knormalE e2
    return $ Let vid ea eb

knormalE (S.UnaryE op e)        = 
    e `normalizeToLet` \v -> return $ Val $ UnaryV op v

knormalE (S.BinE op e1 e2)      = 
    e1 `normalizeToLet` \v1 -> 
    e2 `normalizeToLet` \v2 -> return $ Val $ BinV op v1 v2

-- /Naked/ Rel expressions are translated to Ifte that produces 1 or 0
knormalE (S.RelE op e1 e2)      = 
    e1 `normalizeToLet` \v1 -> 
    e2 `normalizeToLet` \v2 -> 
    return $ Ifte (RelV op v1 v2) (Val $ Int 1) (Val $ Int 0)


knormalE (S.Ifte (S.RelE op e1 e2) e3 e4) =
    e1 `normalizeToLet` \v1 -> 
    e2 `normalizeToLet` \v2 -> do { ea <- knormalE e3
                                  ; eb <- knormalE e4
                                  ; return $ Ifte (RelV op v1 v2) ea eb }

knormalE (S.Ifte e1 e2 e3)      = do
    e1_false <- falseVal e1
    knormalE (S.Ifte (S.RelE EQU e1 e1_false) e2 e3)

knormalE (S.Cmd cmd)            = knormalC cmd
knormalE (e1 S.:>> e2)          = (:>>) <$> knormalE e1 <*> knormalE e2
knormalE (S.Return e)           = Return <$> knormalE e


knormalC :: S.Command -> Mon Expr
knormalC (S.NewPath)                      = return $ Cmd $ NewPath
knormalC (S.ClosePath)                    = return $ Cmd $ ClosePath
knormalC (S.Fill)                         = return $ Cmd $ Fill
knormalC (S.Stroke)                       = return $ Cmd $ Stroke
knormalC (S.MoveTo e1 e2)                 =
    e1 `normalizeToLet` \v1 -> 
    e2 `normalizeToLet` \v2 -> 
    return $ Cmd $ MoveTo v1 v2

knormalC (S.LineTo e1 e2)                 =
    e1 `normalizeToLet` \v1 -> 
    e2 `normalizeToLet` \v2 -> 
    return $ Cmd $ LineTo v1 v2

knormalC (S.CurveTo e1 e2 e3 e4 e5 e6)    = 
    e1 `normalizeToLet` \v1 -> 
    e2 `normalizeToLet` \v2 ->
    e3 `normalizeToLet` \v3 -> 
    e4 `normalizeToLet` \v4 ->
    e5 `normalizeToLet` \v5 -> 
    e6 `normalizeToLet` \v6 ->
    return $ Cmd $ CurveTo v1 v2 v3 v4 v5 v6

knormalC (S.GSave)                        = return $ Cmd $ GSave
knormalC (S.GRestore)                     = return $ Cmd $ GRestore

knormalC (S.Scale e1 e2)                  =
    e1 `normalizeToLet` \v1 -> 
    e2 `normalizeToLet` \v2 -> 
    return $ Cmd $ Scale v1 v2

knormalC (S.Rotate e1)                    =
    e1 `normalizeToLet` \v1 -> 
    return $ Cmd $ Rotate v1

knormalC (S.Translate e1 e2)              =
    e1 `normalizeToLet` \v1 -> 
    e2 `normalizeToLet` \v2 -> 
    return $ Cmd $ Translate v1 v2

knormalC (S.Transform e1 e2 e3 e4 e5 e6)  = 
    e1 `normalizeToLet` \v1 -> 
    e2 `normalizeToLet` \v2 ->
    e3 `normalizeToLet` \v3 -> 
    e4 `normalizeToLet` \v4 ->
    e5 `normalizeToLet` \v5 -> 
    e6 `normalizeToLet` \v6 ->
    return $ Cmd $ Transform v1 v2 v3 v4 v5 v6

knormalC (S.ShowText e1)                  =
    e1 `normalizeToLet` \v1 -> 
    return $ Cmd $ ShowText v1




falseVal :: S.Expr -> Mon S.Expr 
falseVal val = case S.typeOf val of
    TyBool  -> return $ S.Bool False
    TyInt   -> return $ S.Int 0
    TyFloat -> return $ S.Float 0.0
    _       -> throwError $ 
                "falseVal - cannot falsify, type checker missed this." 